/*
 * Copyright (c) Huawei Technologies Co., Ltd. 2022-2022. All rights reserved.
 * Description: alg remote lookup.
 * Author: qingjicheng
 * Create: 2022-11-28
 */

#ifndef REMOTE_LOOK_UP_H
#define REMOTE_LOOK_UP_H
#include "alg_embedding_service.h"

namespace hccl {

class AlgRemoteLookUp : public AlgEmbeddingService {
public:
    explicit AlgRemoteLookUp();
    virtual ~AlgRemoteLookUp() = default;
    HcclResult Process(HcclOpPtr opPtr) override;

protected:
    virtual HcclResult AlgRemoteLookup(HcclOperator<EmbeddingServiceParam> *opPtr);
    virtual HcclResult AlgCollRemoteLookup(HcclOperator<EmbeddingServiceParam> *opPtr);
    virtual HcclResult AlgCollRemoteLookupPaired(HcclOperator<EmbeddingServiceParam> *opPtr);
    virtual HcclResult AlgGetLookupRequest(HcclOperator<EmbeddingServiceParam> *opPtr);
    virtual HcclResult AlgWaitLookupData(HcclOperator<EmbeddingServiceParam> *opPtr);
    virtual HcclResult AlgSetLookupResponse(HcclOperator<EmbeddingServiceParam> *opPtr);
    virtual HcclResult AlgWaitSome(HcclOperator<EmbeddingServiceParam> *opPtr);
    virtual HcclResult AlgKeysDuplicates(HcclOperator<EmbeddingServiceParam> *opPtr);
    virtual HcclResult AlgSendKeys(HcclOperator<EmbeddingServiceParam> *opPtr);
    virtual HcclResult AlgRecvValues(HcclOperator<EmbeddingServiceParam> *opPtr);
    virtual HcclResult AlgRecoverValue(HcclOperator<EmbeddingServiceParam> *opPtr);
    virtual HcclResult AlgResetUniqueHandle(HcclOperator<EmbeddingServiceParam> *opPtr);
    virtual HcclResult AlgWaitSendKeyFinish(HcclOperator<EmbeddingServiceParam> *opPtr);
    HcclResult AlgResetUniqueTask(HcclOperator<EmbeddingServiceParam> *opPtr);

private:

    FuncTable funcTable_;

    void AlgCollRemoteLookupParaAssign(const HcclOperator<EmbeddingServiceParam> &opPtr,
        PartitionMapTask &partitionMapTask);
};
}
#endif
